package edu.kit.csl.pisa.models;

/*
This file is part of the PISA Alignment Tool.

Copyright (C) 2013
Karlsruhe Institute of Technology
Cognitive Systems Lab (CSL)
Felix Stahlberg

PISA is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

PISA is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with PISA. If not, see <http://www.gnu.org/licenses/>.
*/

import java.util.Arrays;

import edu.kit.csl.pisa.datatypes.AlignmentPosition;
import edu.kit.csl.pisa.datatypes.AlignmentSourceWord;
import edu.kit.csl.pisa.datatypes.BadParameterValueException;
import edu.kit.csl.pisa.datatypes.Corpus;
import edu.kit.csl.pisa.datatypes.Dictionary;
import edu.kit.csl.pisa.datatypes.SentenceAlignment;
import edu.kit.csl.pisa.io.Logger;

/**
 * This class implements the Model 4-phone alignment model. Model 4P uses a 
 * distortion model oriented on IBM-4 instead of IBM-3 and adds a src/trgt
 * sentence length mismatch parameter l, which is distributed according a
 * sentence length model. In-depth discussion this model can be found in the 
 * technical paper. The parameters are stored as n-dimensional double arrays. 
 * The order in which the dimensions are assigned is as follows:<br>
 * At first, the parameters on the right side of the "|" in the order in which 
 * these dependencies are denoted in the technical paper. The last dimension
 * will always store the left side of the "|" (in P(.|.) ):<br>
 * o:	1. English Word (e_A_0j)<br>
 * 		2. Relative French word position (A_2j)<br>
 * 		3. French word length in phonemes (A_3j)<br>
 * d1:  1. even: positive values, uneven: negative
 * d2:  1. relative position to last word in cept
 */
public class Model4P extends ModelP {
	
	/*
	 * Model parameter. Which dimension contains which values is denoted in the
	 * head of this class.
	 */
	protected double[][] o;
	// even indexes store positive values, uneven negative values
	protected double[] d1;
	protected double[] d2;
	
	protected double senLenWeight = 1.0f;
	
	// Configuration parameters
	
	// Model instance for collecting fractional counts
	private Model4P newModel;
	
	private final int maxDistortion;
	private final int maxDistortion2;
	private SentenceLengthModel senLenModel;
	
	/**
	 * Sole constructor.
	 * 
	 * @param name name of the model
	 * @throws BadParameterValueException if the sentence length model is not
	 * 		known
	 */
	public Model4P(String name) throws BadParameterValueException {
		super(name);
		maxDistortion = getConfigInteger("MaxDistortion");
		maxDistortion2 = maxDistortion * 2;
		String str = getConfigString("SentenceLengthModel");
		senLenWeight = getConfigFloat("SentenceLengthModelWeight");
		if (str.equals("difference")) {
			senLenModel = new DifferenceSentenceLengthModel(this,
					maxDistortion);
		} else if (str.equals("uniform")) {
			senLenModel = new UniformSentenceLengthModel();
		} else if (str.equals("negativeBinomial")) {
			senLenModel = new NBSentenceLengthModel(this);
		} else if (str.equals("restrictedNegativeBinomial")) {
			senLenModel = new RestrictedNBSentenceLengthModel(this);
		} else {
			throw new BadParameterValueException("Sentence length model '"
				+ str + "' not implemented.");
		}
	}

	/**
	 * Calculate the log probability P(a,f|e). The formula can be found in the
	 * technical paper.  
	 * 
	 * @return the logarithm of P(a,f|e)
	 */
	@Override
	public double calculateAlignmentProbability(SentenceAlignment a) {
		// TODO: Model 3P: No o for NULL?
		final int srcLength = a.getSourceSentence().getLength(); // l
		final int nTrgtWords = a.getAlignment().countTargetWords(); // k
		double[] summands = new double[4 + 2 * srcLength + 2 * nTrgtWords
		                               + a.getTargetSentence().getLength()];
		int[] ceptSum = new int[srcLength];
		int[] ceptCount = new int[srcLength];
		short[] ceptLastTrgtPos = new short[srcLength];
		// p0, p1, fertility, lexical translation
		int i = calculateModelPAlignmentProbability(summands, a);
		// Word length and distortion
		for (AlignmentPosition pos : a.eachTargetWord()) {
			int srcPos = pos.getSourceWordPosition() - 1;
			short trgtPos = pos.getTargetWordPosition();
			if (srcPos == -1) {
				continue;
			}
			if (ceptCount[srcPos]++ > 0) {
				summands[i++] = d2[Math.min(maxDistortion,
						(int) (trgtPos - ceptLastTrgtPos[srcPos]))];
			}
			ceptLastTrgtPos[srcPos] = trgtPos;
			ceptCount[srcPos] += trgtPos;
			summands[i++] = o[pos.getSourceWord()]
							 [pos.getTargetWordLength()];
		}
		for (AlignmentSourceWord srcWord : a.eachSourceWordWithoutNull()) {
			final short k = srcWord.getSrcWordPos();
			if (ceptCount[k] > 0) {
				int dist = 2 * (k - ceptSum[k] / ceptCount[k]);
				if (dist < 0) {
					dist = dist * (-1) - 1;
				}
				summands[i++] = d1
						[getD1Index(k, ceptSum[k] / ceptCount[k])];
			}
		}
		summands[i++] = 
			senLenWeight * senLenModel.getProbability(srcLength, nTrgtWords);
		return sum(summands, 0, i);
	}


	/**
	 * t, n, d are written to the file system.
	 * {@inheritDoc}
	 */
	@Override
	public void dumpToFilesystem(String prefix, String postfix) {
		Dictionary srcDict = Dictionary.getSingleton(Dictionary.SRC_DICT);
		Dictionary trgtDict = Dictionary.getSingleton(Dictionary.TRGT_DICT);
		this.dumpParameter(prefix + "n" + postfix, n);
		this.dumpParameter(prefix + "actual.n" + postfix, n, new Dictionary[]{
			srcDict, null});
		this.dumpParameter(prefix + "t" + postfix, t);
		this.dumpParameter(prefix + "actual.t" + postfix, t, new Dictionary[]{
			srcDict, null, trgtDict});
		this.dumpParameter(prefix + "o" + postfix, o);
		this.dumpParameter(prefix + "actual.o" + postfix, o, new Dictionary[]{
			srcDict, null, null});
		double[][] arr = new double[1][];
		arr[0] = d1;
		this.dumpParameter(prefix + "d1" + postfix, arr);
		arr[0] = d2;
		this.dumpParameter(prefix + "d2" + postfix, arr);
		senLenModel.dumpToFilesystem(prefix, postfix);
	}

	/* (non-Javadoc)
	 * @see edu.kit.csl.pisa.models.AlignmentModel#initializeFractionalCounts()
	 */
	@Override
	public void initializeFractionalCounts() {
		try {
			newModel = new Model4P(getName());
		} catch (BadParameterValueException e) {
			; // Already handled in first place
		}
		newModel.allocateModelPMemory(false);
		newModel.allocateMemory(false);
	}

	/* (non-Javadoc)
	 * @see edu.kit.csl.pisa.models.AlignmentModel#writeBackFractionalCounts()
	 */
	@Override
	public void writeBackFractionalCounts() {
		normalizeParameters();
		this.n = newModel.n;
		this.t = newModel.t;
		this.o = newModel.o;
		this.d1 = newModel.d1;
		this.d2 = newModel.d2;
		this.p0 = newModel.p0;
		this.p1 = newModel.p1;
		this.senLenModel = newModel.senLenModel;
		newModel = null;
	}
	
	/*
	 * This is a helper method of {@link #writeBackFractionalCounts()}.
	 * Normalize parameters in a separate model instance. Thereby we assume,
	 * that the parameters are containing fractional counts. This method
	 * realizes a limited amount of parallelism, in which each parameter is
	 * normalized in a separate thread.
	 */
	private void normalizeParameters() {
		final int nWorkers = 5;
		Thread[] workers = new Thread[nWorkers];
		workers[0] = new Thread() {
			public void run() {
				for (double[] a1 : newModel.n) {
					normalize(a1);
				}
			}
		};
		workers[1] = new Thread() {
			public void run() {
				for (double[] a1 : newModel.o) {
					normalize(a1);
				}
			}
		};
		workers[2] = new Thread() {
			public void run() {
				for (double[][] a1 : newModel.t) {
					for (double[] a2 : a1) {
						normalize(a2);
					}
				}
			}
		};
		workers[3] = new Thread() {
			public void run() {
				double[] a1 = newModel.d1;
				//for (double[] a1 : newModel.d1) {
					normalize(a1);
				//}
			}
		};
		workers[4] = new Thread() {
			public void run() {
				double[] a1 = newModel.d2;
				//for (double[] a1 : newModel.d2) {
					normalize(a1);
				//}
			}
		};
		for (int i = 0; i < nWorkers; i++) {
			workers[i].start();
		}
		newModel.senLenModel.normalize();
		for (int i = 0; i < nWorkers; i++) {
			try {
				workers[i].join();
			} catch (InterruptedException e) {
				Logger.getSingleton().error("Join of worker threads in EM "
						+ "training interrupted: " + e.getMessage());
			}
		}
	}

	/**
	 * This template method (GoF) is the interface for implementing subclasses
	 * to multi-threaded collection of fractional counts.
	 * TODO: Use recursive D&C summation
	 * 
	 * @param pair sentence alignment
	 * @param weight weight of the alignment
	 * @see #collectFractionalCounts(Corpus)
	 */
	@Override
	public void aggregateFractionalCount(SentenceAlignment a, double weight) {
		aggregateModelPFractionalCount(newModel, a, weight);
		synchronized (newModel.o) {
			for (AlignmentPosition pos : a.eachTargetWord()) {
				if (pos.getSourceWordPosition() == 0) {
					continue;
				}
				newModel.o	[pos.getSourceWord()]
							[pos.getTargetWordLength()] += weight;
			}
		}
		final int srcLength = a.getSourceSentence().getLength();
		int[] ceptSum = new int[srcLength];
		int[] ceptCount = new int[srcLength];
		short[] ceptLastTrgtPos = new short[srcLength];
		synchronized (newModel.senLenModel) {
			newModel.senLenModel.count(weight, srcLength, 
					a.getAlignment().countTargetWords());
		}
		synchronized (newModel.d1) {
			for (AlignmentPosition pos : a.eachTargetWord()) {
				int srcPos = pos.getSourceWordPosition() - 1;
				short trgtPos = pos.getTargetWordPosition();
				if (srcPos == -1) {
					continue;
				}
				if (ceptCount[srcPos]++ > 0) {
					newModel.d2[Math.min(maxDistortion, 
							trgtPos - ceptLastTrgtPos[srcPos])] += weight;
				}
				ceptLastTrgtPos[srcPos] = trgtPos;
				ceptCount[srcPos] += trgtPos;
			}
			for (AlignmentSourceWord srcWord : a.eachSourceWordWithoutNull()) {
				final short k = srcWord.getSrcWordPos();
				if (ceptCount[k] > 0) {
					newModel.d1
							[getD1Index(k, ceptSum[k] / ceptCount[k])] += weight;
				}
			}
		}
	}
	
	/**
	 * Get index in d1 from source and cept position. Distortions are stored
	 * in this order: 0,-1,1,-2,2,-3,3...
	 * 
	 * @param srcPos source word position
	 * @param ceptPos position of cept in the target sequence
	 */
	protected int getD1Index(int srcPos, int ceptPos) {
		int dist = 2 * (srcPos - ceptPos);
		return Math.min(maxDistortion2, (dist < 0) 
				? (-dist - 1) : dist);
	}

	/**
	 * Import model parameters from ModelP ({@link ModelP}). If the given 
	 * instance is a {@link Model3P}, the word length model is also imported,
	 * and the distortion model is discarded. Note that the parameter arrays 
	 * are NOT copied to avoid additional memory overhead and time to copy. 
	 * However, keep in mind that this could cause side effects.
	 * 
	 * @param srcModel should be a {@link Model3P} instance
	 * @throws IllegalArgumentException if srcModel is not a 
	 * 				{@link Model3P} instance
	 */
	@Override
	public void importModelParameters(AlignmentModel srcModel)
			throws IllegalArgumentException {
		if (srcModel instanceof ModelP) {
			allocateMemory(true);
			ModelP modelP = (ModelP) srcModel;
			this.n = modelP.n;
			this.t = modelP.t;
			this.p0 = modelP.p0;
			this.p1 = modelP.p1;
			if (srcModel instanceof Model3P) {
				this.o = ((Model3P) srcModel).getO();
			}
			// Uniformly distribute d1,2
			double prop = Math.log(1f / maxDistortion2);
			//for (double[] a1 : d1) {
				Arrays.fill(d1, prop);
			//}
			prop = Math.log(1f / maxDistortion);
			//for (double[] a1 : d2) {
				Arrays.fill(d2, prop);
			//}
		} else {
			throw new IllegalArgumentException("Import is not implemented"
				+ " for this alignment model.");
		}
	}
	
	/*
	 * Allocate memory for model parameters.
	 * 
	 * @param print Set to TRUE to print parameter sizes on the screen
	 */
	private void allocateMemory(boolean print) {
		final int srcDictSize = 
				Dictionary.getSingleton(Dictionary.SRC_DICT).size() + 1;
		Logger log = Logger.getSingleton();
		// Allocate memory
		if (print) {
			log.debug("Model-4P o parameter: " + AlignmentModel
					.humanReadableDoubleCount(srcDictSize * maxWordLength));
			log.debug("Model-4P d1,2 parameter: " + AlignmentModel
					.humanReadableDoubleCount(srcDictSize * 3));
		}
		o = new double[srcDictSize][maxWordLength];
		d1 = new double[maxDistortion2 + 1];
		d2 = new double[maxDistortion + 1];
	}
}
